import os

import numpy as np
import cv2
import torch
import torchvision

from dataset import Crack

class MyFCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.weights_of_fcn = torchvision.models.segmentation.FCN_ResNet50_Weights.DEFAULT  # 默认权重
        self.FCN = torchvision.models.segmentation.fcn_resnet50(weights = self.weights_of_fcn)
        self.FCN.classifier[4] = torch.nn.Conv2d(in_channels=512, out_channels=2, kernel_size=(1,1), stride=(1,1))
        self.FCN.aux_classifier[4] = torch.nn.Conv2d(in_channels=256, out_channels=2, kernel_size=(1,1), stride=(1,1))

    def forward(self, imgs):
        imgs = self.FCN(imgs)
        return imgs




if __name__ == '__main__':
    model = MyFCN()
    dataset = Crack(r'./data/train/imgs', r'./data/train/masks')
    dataset_loader = torch.utils.data.DataLoader(dataset = dataset, batch_size=4, shuffle=True)

    model.eval()
    for imgs, labels in dataset_loader:
        result = model(imgs)
        print(result['out'].shape)
        break


